"""
File: 国家数据库爬虫例子
Data: 2024/4/27
"""
import requests
import time
import json
import pandas as pd

BASE_URL = "https://data.stats.gov.cn"
HEADERS = {
    "accept": "application/json, text/javascript, */*; q=0.01",
    "accept-language": "zh-CN,zh;q=0.9,fr;q=0.8,en;q=0.7",
    "sec-ch-ua": "\"Chromium\";v=\"124\", \"Google Chrome\";v=\"124\", \"Not-A.Brand\";v=\"99\"",
    "sec-ch-ua-mobile": "?0",
    "sec-ch-ua-platform": "\"Windows\"",
    "sec-fetch-dest": "empty",
    "sec-fetch-mode": "cors",
    "sec-fetch-site": "same-origin",
    "x-requested-with": "XMLHttpRequest",
    "cookie": "wzws_sessionid=gmZjNWVlMaBmLJpSgDIyMy43My4xNDAuMjQzgWExNmRkYQ==; u=2; acmrAutoLoginUser=null; acmrAutoSessionId=null; JSESSIONID=0i4elkzz82FNoztXfOVWNvkvRSbSH2iH08KZHpdr1g71jB-VfxCu!460167158",
    "Referer": "https://data.stats.gov.cn/easyquery.htm?cn=E0103",
    "Referrer-Policy": "strict-origin-when-cross-origin"
}


# 获取爬虫列表
def get_url_list(params):
    """
    固定资产投资和房地产的详细数据列表
    :return:
    """
    url = f"{BASE_URL}/easyquery.htm"
    # 请求体
    params["m"] = "getTree"

    # 发送POST请求
    response = requests.post(url, data=params, verify=False)
    return response.json()


def get_data(item):
    """

    :param item: 从列表返回的参数，填入参数
    :return:
    """
    url = f"{BASE_URL}/easyquery.htm"
    params = {
        'm': 'QueryData',
        'dbcode': item['dbcode'],
        'rowcode': item['wdcode'],
        'colcode': 'reg',
        'wds': json.dumps([{"wdcode": "sj", "valuecode": item["year"]}]),
        'dfwds': json.dumps([{"wdcode": item["wdcode"], "valuecode":item["id"]}]),
        'k1': str(int(time.time() * 1000)),
        'h': '1'
    }
    # 发送GET请求
    response = requests.get(url, params=params, verify=False)
    data = response.json()
    # 返回数据
    # datanodes: 数据， wdnodes: 数据说明
    return data["returndata"]["datanodes"], data["returndata"]["wdnodes"]


def save_data(data, wdnode):
    """
    把数据保存为csv，或按要求变成其他类型
    :param data: 数据
    :param wdnode: 数据说明
    :return:
    """
    insert_data = []
    zb_dict, reg_dict = handle_data_to_dict(wdnode)
    for item in data:
        tmp_data = []
        for wd in item["wds"]:
            if wd["wdcode"] == "zb":
                tmp_data.append(zb_dict[wd["valuecode"]])
            elif wd["wdcode"] == "reg":
                tmp_data.append(reg_dict[wd["valuecode"]])
        tmp_data.append(item["data"]["strdata"])
        insert_data.append(tmp_data)
    
    df = pd.DataFrame(insert_data)
    print(df.head())


def handle_data_to_dict(wdnode):
    """
    处理数据字典
    :param wdnode:
    :return:
    """
    zb_dict = {}
    reg_dict = {}
    for item in wdnode:
        if item["wdcode"] == "zb":
            for node in item["nodes"]:
                # print(node["cname"], node["code"], node["unit"])
                zb_dict[node["code"]] = f'{node["cname"]}({node["unit"]})'
        if item["wdcode"] == "reg":
            for node in item["nodes"]:
                # print(node["cname"], node["code"])
                reg_dict[node["code"]] = node["cname"]
    return zb_dict, reg_dict


if __name__ == '__main__':
    # 数据列表参数
    data_config = {
        "id": "A05",
        "dbcode": "fsnd",
        "wdcode": "zb"
    }
    items = get_url_list(data_config)
    # 请求数据年份
    years = range(2022, 2018, -1)
    for item in items:
        for year in years:
            print(f"请求{item['name']}数据")
            item["year"] = str(year)
            data, wdnode = get_data(item)
            print(data[:3])
            print(wdnode)
            save_data(data, wdnode)
            break
        break
